'''
Author: boboji11 wendychen112@qq.com
Date: 2025-06-21 18:14:14
LastEditors: boboji11 wendychen112@qq.com
LastEditTime: 2025-06-21 18:15:14
FilePath: \预选赛代码\code.py
Description:
Copyright (c) 2025 by boboji11 , All Rights Reserved. 
'''
from ultralytics import YOLO
import cv2
import time
# record time
nowTime=time.time()
lastTime=nowTime
# load modal
yoloPath = "best.engine" # path of the weight of modal
yolo11 = YOLO(yoloPath, task='classify') # set task as 'classify'

# NOTE: 4 th camera is realsense d435i RGB camera in my machine
# please comfirm the serial number of your machine of D435i
cap = cv2.VideoCapture(4)
if not cap.isOpened():
    print("Error: Could not open camera.please comfirm the serial number of your machine of D435i for OpenCV")
    exit()
# store the result of the prediction
predict_list = []
class_dict={0: '偏高', 1: '偏低', 2: '正常'}
class_dict_predict = {0: 0, 1:0, 2:0}
print("yolo prediction start!")
# Predcition Loop
while True:
    ret, frame = cap.read()
    if not ret:
        print("Error: Could not read frame.")
        continue
    frame = frame[120:360, 160:480] # clip
    frame = cv2.resize(frame, (640, 480),interpolation=cv2.INTER_LINEAR) # extend
    results = yolo11.predict(frame, imgsz=224, verbose=False) # predict
    nowTime = time.time()
    # analysis the results
    for result in results:
        prob_top1 = result.probs.top1conf.item()
        # only save the results whose confidence probability is more than 0.9
        if prob_top1 < 0.9:
            continue
        lastTime=nowTime
        num_class_top1 = result.probs.top1
        #class_top1 = class_dict[num_class_top1]
        predict_list.append((num_class_top1, prob_top1))
    if nowTime - lastTime >2
        # clear the prediction list if there has been no new useful predition result for 2s
        predict_list.clear()
        lastTime=nowTime
    if len(predict_list) >= 10:
        # statistics
        for item in predict_list:
            class_dict_predict[item[0]] += 1
        # choose the most possible
        max_class = max(class_dict_predict, key=class_dict_predict.get)
        # clear
        for name, value in class_dict_predict.items():
            class_dict_predict[name] = 0
        predict_list.clear()
        # print the results
        print("预测结果是",class_dict[max_class])
cap.release()